import sys
import os
sys.path.insert(0, os.path.abspath(os.path.join(os.path.dirname(__file__), "..")))

import numpy as np
from ml_lib.core import Tensor
from ml_lib.nn.layers import Linear, Sequential
from ml_lib.nn.loss import MSELoss
from ml_lib.optim import SGD

def test_simple_single_task():
    """测试单任务学习是否正常"""
    print("=" * 50)
    print("测试单任务学习：y = 2x + 1")
    print("=" * 50)
    
    # 创建简单模型
    model = Sequential(
        Linear(1, 10),
        Linear(10, 1)
    )
    
    optimizer = SGD(model.parameters(), lr=0.01)
    criterion = MSELoss()
    
    # 创建数据
    x_data = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32)
    y_data = 2 * x_data + 1  # y = 2x + 1
    
    print("训练数据:")
    for i in range(len(x_data)):
        print(f"x={x_data[i][0]:.1f}, y={y_data[i][0]:.1f}")
    
    # 训练
    for epoch in range(1000):
        total_loss = 0.0
        for i in range(len(x_data)):
            x = Tensor(x_data[i:i+1], requires_grad=True)
            target = Tensor(y_data[i:i+1])
            
            optimizer.zero_grad()
            output = model(x)
            loss = criterion(output, target)
            loss.backward()
            optimizer.step()
            
            total_loss += float(loss.data)
        
        if epoch % 20 == 0:
            print(f"Epoch {epoch}, Loss: {total_loss/len(x_data):.6f}")
    
    # 测试预测
    print(f"\n最终预测:")
    for i in range(len(x_data)):
        x = Tensor(x_data[i:i+1])
        pred = model(x)
        print(f"x={x_data[i][0]:.1f}, 预测={float(pred.data):.3f}, 真实={y_data[i][0]:.1f}")
    
    # 测试新数据
    test_x = Tensor(np.array([[5.0]], dtype=np.float32))
    test_pred = model(test_x)
    print(f"测试: x=5.0, 预测={float(test_pred.data):.3f}, 期望=11.0")


def test_manual_mtl():
    """手动实现多任务学习"""
    print("\n" + "=" * 50)
    print("手动多任务学习：y1 = 2x + 1, y2 = x²")
    print("=" * 50)
    
    # 共享编码器
    shared_encoder = Sequential(
        Linear(1, 10)
    )
    
    # 任务特定头部
    head1 = Linear(10, 1)  # 线性任务
    head2 = Linear(10, 1)  # 二次任务
    
    # 优化器
    all_params = list(shared_encoder.parameters()) + list(head1.parameters()) + list(head2.parameters())
    optimizer = SGD(all_params, lr=0.01)
    criterion = MSELoss()
    
    # 数据
    # x_data = np.array([[1.0], [2.0], [3.0], [4.0]], dtype=np.float32)
    x_data = np.linspace(0, 4, 100,dtype=np.float32).reshape(-1,1)
    # x_data = x_data.reshape(-1, 1)
    y1_data = 2 * x_data + 1  # 线性
    y2_data = x_data ** 2     # 二次
    
    print("训练数据:")
    for i in range(len(x_data)):
        print(f"x={x_data[i][0]:.1f}, y1={y1_data[i][0]:.1f}, y2={y2_data[i][0]:.1f}")
    
    # 训练
    for epoch in range(2000):
        total_loss1 = 0.0
        total_loss2 = 0.0
        
        for i in range(len(x_data)):
            x = Tensor(x_data[i:i+1], requires_grad=True)
            target1 = Tensor(y1_data[i:i+1])
            target2 = Tensor(y2_data[i:i+1])
            
            optimizer.zero_grad()
            
            # 前向传播
            shared_features = shared_encoder(x)
            pred1 = head1(shared_features)
            pred2 = head2(shared_features)
            
            # 计算损失
            loss1 = criterion(pred1, target1)
            loss2 = criterion(pred2, target2)
            total_loss = loss1 + loss2  # 等权重
            
            # 反向传播
            total_loss.backward()
            optimizer.step()
            
            total_loss1 += float(loss1.data)
            total_loss2 += float(loss2.data)
        
        if epoch % 50 == 0:
            print(f"Epoch {epoch}, Loss1: {total_loss1/len(x_data):.6f}, Loss2: {total_loss2/len(x_data):.6f}")
    
    # 测试预测
    print(f"\n最终预测:")
    for i in range(len(x_data)):
        x = Tensor(x_data[i:i+1])
        shared_features = shared_encoder(x)
        pred1 = head1(shared_features)
        pred2 = head2(shared_features)
        print(f"x={x_data[i][0]:.1f}, pred1={float(pred1.data):.3f} (真实{y1_data[i][0]:.1f}), pred2={float(pred2.data):.3f} (真实{y2_data[i][0]:.1f})")


if __name__ == "__main__":
    # test_simple_single_task()
    test_manual_mtl() 